#所需依赖
import os
import sys
root_path=os.path.abspath(os.path.join(os.path.dirname(__file__),"."))
sys.path.append(root_path)

import cv2
import numpy as np
from utility import Init_Paddle
from general import check_img_size,letterbox,non_max_suppression_paddle,scale_coords,Annotator,colors
from pathlib import Path
import time
import glob

debug=False

#数据预处理
def preprocess(img,imgsz):
    '''
    :param: img:图片Mat矩阵
    :param: imgsz: 期望得到的图片较长边尺寸，默认为640
    '''
    # 查看图片尺寸能被stride整除的最大尺寸
    imgsz = check_img_size(imgsz, s=64)
    # 等比例缩放
    img=letterbox(img,imgsz,64,auto=False)[0]
    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    img = np.ascontiguousarray(img) #将一个内存不连续存储的数组转换为内存连续存储的数组，使得运行速度更快。
    img = img.astype(np.float32)
    img /= 255.0  # 归一化
    img=img[np.newaxis,:]
    return img

class Paddle_Infer():
    def __init__(self):
        model_dir = "yolov5/white_model/inference_model"  # 模型路径
        self.predictor, self.input_tensor, self.output_tensors, self.config = Init_Paddle(model_dir,use_tensorrt=False)
        self.classes = []  # 类别信息
        self.save_path = root_path + "/runs/inference_results"  # 推理结果保存路径
        if debug:
            os.makedirs(self.save_path, exist_ok=True)

    def predict_img(self, img, imgpath, imgsz=640):
        '''
        :param img:图片Mat矩阵
        :param imgsz: 期望得到的缩放后的图片尺寸
        '''
        # im0=img.copy() #原先的图片
        img_pre = preprocess(img, imgsz)  # resize后的图片
        self.input_tensor.copy_from_cpu(img_pre)

        self.predictor.run()
        output = self.output_tensors[0].copy_to_cpu()
        pred = non_max_suppression_paddle(output, 0.25, 0.5, None, False, max_det=1000)
        
        # 释放内存池中的所有临时 Tensor
        self.predictor.try_shrink_memory()

        # 可视化
        x, j = self.vis_show(pred, img, img_pre, imgpath)
        return x, j

    def vis_show(self, pred, image, image_padding, img_path):
        '''
        :param: pred: 预测结果:
        :param: image_old:原始图像
        :param: image_padding: 缩放后的图片
        :param: imgpath: 图片路径
        '''
        # names=[f'class{i}'for i in range(1000)]
        # names = self.classes

        names = ['white_dolphin']


        # print(names)
        flag = 0
        for i, det in enumerate(pred):
            # 绘制预测框
            annotator = Annotator(image, line_width=3, example=str(names))
            # 检测到目标时
            if len(det):
                flag = 1
                det[:, :4] = scale_coords(image_padding.shape[2:], det[:, :4], image.shape).round()
                # 写结果
                for *xyxy, conf, cls in reversed(det):
                    c = int(cls)
                    label = f'{names[c]} {conf:.2f}'
                    annotator.box_label(xyxy, label, color=colors(c, True))

            im0 = annotator.result()

            # 保存推理结果
            if debug:
                imgname = Path(img_path).name  # 获取文件名
                cv2.imwrite(self.save_path + os.sep + imgname, im0)
            return im0, flag


def main(img):
    s = time.time()
    pre = Paddle_Infer()
    path = './test.jpg'
    res, j = pre.predict_img(img, path)
    e = time.time()
    print("yolov5_predict_time : {}s".format(e - s))
    if j == 0:
        return [img, 0]
    else:
        return [img, 1]
